import os
import joblib
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, cross_val_score, StratifiedKFold
from sklearn.metrics import (accuracy_score, classification_report, 
                            confusion_matrix, roc_curve, roc_auc_score)
from sklearn.preprocessing import LabelEncoder

# 设置中文字体 [^优化]
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

def build_classification_model(feature_df, save_path=''):
    """分类模型构建（性能优化）[10](@ref)"""
    # 1. 数据准备
    X = feature_df.drop(['Group', 'Subject_ID'], axis=1, errors='ignore')
    y = LabelEncoder().fit_transform(feature_df['Group'])
    
    # 2. 分层分割（保持类别比例）[4](@ref)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, stratify=y, random_state=42
    )
    
    # 3. 训练随机森林（参数优化）[10](@ref)
    model = RandomForestClassifier(
        n_estimators=150,           # 减少树数量加速训练
        max_depth=7,                # 限制深度防止过拟合
        class_weight='balanced',
        n_jobs=-1,                  # 并行加速 [11](@ref)
        random_state=42
    )
    model.fit(X_train, y_train)
    
    # 4. 评估模型
    y_pred = model.predict(X_test)
    y_proba = model.predict_proba(X_test)[:,1]
    
    # 核心指标计算
    acc = accuracy_score(y_test, y_pred)
    auc = roc_auc_score(y_test, y_proba)
    report = classification_report(y_test, y_pred, target_names=['TD', 'ASD'])
    
    # 5. 交叉验证（分层K折）[4](@ref)
    cv = StratifiedKFold(n_splits=5)
    cv_scores = cross_val_score(model, X, y, cv=cv, scoring='accuracy')
    
    # 6. 保存结果（UTF-8编码）[6](@ref)
    if save_path:
        os.makedirs(save_path, exist_ok=True)
        joblib.dump(model, os.path.join(save_path, 'asd_classifier.pkl'))
        
        feat_importance = pd.DataFrame({
            'Feature': X.columns,
            'Importance': model.feature_importances_
        }).sort_values('Importance', ascending=False)
        feat_importance.to_csv(
            os.path.join(save_path, 'feature_importance.csv'), 
            index=False, encoding='utf-8-sig'
        )
        
        with open(os.path.join(save_path, 'classification_report.txt'), 'w', encoding='utf-8') as f:
            f.write(f"准确率: {acc:.4f}\n")
            f.write(f"AUC值: {auc:.4f}\n")
            f.write(f"交叉验证准确率: {np.mean(cv_scores):.4f}±{np.std(cv_scores):.4f}\n\n")
            f.write(report)
    
    # 7. 可视化结果
    plot_model_performance(y_test, y_proba, save_path)
    plot_feature_importance(feat_importance, save_path)
    
    return model, report

def plot_model_performance(y_true, y_proba, save_path):
    """模型性能图（中文标签）[^优化]"""
    plt.figure(figsize=(15, 5))
    
    # 1. 混淆矩阵
    plt.subplot(131)
    cm = confusion_matrix(y_true, y_proba > 0.5)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=['TD', 'ASD'], yticklabels=['TD', 'ASD'])
    plt.title('混淆矩阵')
    plt.ylabel('真实标签')
    plt.xlabel('预测标签')
    
    # 2. ROC曲线
    plt.subplot(132)
    fpr, tpr, _ = roc_curve(y_true, y_proba)
    plt.plot(fpr, tpr, label=f'AUC = {roc_auc_score(y_true, y_proba):.2f}')
    plt.plot([0, 1], [0, 1], 'k--')
    plt.title('ROC曲线')
    plt.xlabel('假阳性率')
    plt.ylabel('真阳性率')
    plt.legend()
    
    # 3. 概率分布
    plt.subplot(133)
    for label, name in zip([0, 1], ['TD', 'ASD']):
        sns.kdeplot(y_proba[y_true == label], label=name, fill=True)
    plt.title('预测概率分布')
    plt.xlabel('预测概率(ASD)')
    plt.ylabel('密度')
    plt.legend()
    
    plt.tight_layout()
    if save_path:
        plt.savefig(os.path.join(save_path, 'model_performance.png'), 
                   dpi=300, bbox_inches='tight')
    plt.close()

def plot_feature_importance(feat_importance, save_path, top_n=10):
    """特征重要性（中文标签）[^优化]"""
    plt.figure(figsize=(10, 6))
    top_features = feat_importance.head(top_n)
    sns.barplot(x='Importance', y='Feature', data=top_features, palette='viridis')
    plt.title(f'Top {top_n} 重要特征')
    plt.tight_layout()
    if save_path:
        plt.savefig(os.path.join(save_path, 'feature_importance.png'), 
                   dpi=300, bbox_inches='tight')
    plt.close()